#!/usr/bin/perl

# 指定物种名或物种分类名，提取该物种下面全部的蛋白序列名称

# 附属文件的制作方法：
# gunzip -c prot.accession2taxid.gz|cut -f 1,3|gzip -c >NR2Tax.txt.gz
# zgrep "NCBI_TaxID" idmapping.dat.2015_03.gz |cut -f 1,3|gzip -c >UniProt2Tax.txt.gz
# perl NodeExtract.pl |gzip -c >TaxInfo.txt.gz

use strict;
use warnings;
use Getopt::Long;
use Cwd;
use Cwd 'abs_path'; #用abs_path函数可将相对路径变成绝对路径
use Switch;
use FindBin qw($Bin);
use Parallel::Simple qw(prun);
use File::Basename;

our $TaxInfo="$Bin/TaxInfo.txt.gz"; # Taxonomy物种分类层级详细信息
our $NRInfo="$Bin/NR2Tax.txt.gz"; # NR库对蛋白的总表
our $UniProtInfo="$Bin/UniProt2Tax.txt.gz"; #UniProt对Taxinomy的总表

my $TaxID;
my $Name;
my $list;
my $out;
my $prefix;
my $help;

GetOptions(
	  'taxid=i' => \$TaxID,
	  'list=s' => \$list,
	  'name=s' => \$Name,
	  'prefix=s' => \$prefix,
	  'outdir=s' => \$out,
	  'h|help' => \$help,
);

my $usage=<<USAGE;
perl $0 
	给定物种或物种大类的拉丁学名或TaxID，挑选出NR库和UniProt库中的蛋白ID
	-taxid 			NCBI的Taxonomy ID号，推荐输入这个
	-name 		物种拉丁学名。若双名法中有空格需要外面加引号。输入该名称可能找不到必要的物种
	-list 		首列为taxid的物种名称列表，可以由本程序生成，也可以自己制作。但是，该列表中的taxid不再进一步寻找低级的分类条目了
	-outdir			输出文件夹，默认为当前文件夹
	-prefix			文件输出结果前缀，默认物种拉丁名（名称中的空格自动替换为下划线）
	-h/help			show this message
	
USAGE

if($help){print $usage; exit();}
unless($Name or $TaxID or $list){print $usage; exit();}

unless(-e $out){system("mkdir -p $out");}

my %SelectedTax=(); # 选中的物种的哈希表

unless($list){ #未输入物种表格
	if($Name and $TaxID){
		print STDERR "物种学名和TaxID号只要输一个就好啦^_^ \n";
		exit();
	}

	my ($TaxName,$Tax)=&LoadTaxInfo(); #载入物种信息
	my %TaxName=%{$TaxName};
	my %Tax=%{$Tax};
	if($Name){
		$TaxID=$TaxName{$Name};
		unless($TaxID){
			print STDERR "找不到物种名$Name \n";
			exit(); #输入的名称无效，需要用户检查
		}
	}

	unless($prefix){
		$prefix=$TaxName{$TaxID};
		$prefix=~s/\ /\_/g; #空格全部修改成短下划线
	}

	# 测试区在这里
	%SelectedTax=&GetFullTaxID($TaxID,\%Tax,\%TaxName);
	open O,">$out/$prefix.taxid.xls";
	foreach (keys %SelectedTax){ #先把TaxID打印出来再说
		print O "$_\t$SelectedTax{$_}\n"; # GetFullTaxID中已经把TaxName加了，这里不必再加
	}
	close O; 
	# 测试区结束
}
else{
	open I,"$list";
	while(<I>){
		chomp;
		my @a=split/\t/;
		$_=shift(@a);
		$SelectedTax{$_}=join("\t",@a);
	}
	close I;
	unless($prefix){
		$prefix=(split/\//,$list)[-1]; #未指定prefix就直接用输入文件名代替吧
	}
}

prun(
	sub1=>[\&DB2Tax,(\%SelectedTax,$UniProtInfo,"$out/$prefix.UniProt.txt","$out/$prefix.UniProt.stat.xls")],
	sub2=>[\&DB2Tax,(\%SelectedTax,$NRInfo,"$out/$prefix.NR.txt","$out/$prefix.NR.stat.xls")],
)or die(Parallel::Simple::errplus());

# &DB2Tax(\%SelectedTax,$UniProtInfo,"$out/$prefix.UniProt.txt","$out/$prefix.UniProt.stat.xls");
# &DB2Tax(\%SelectedTax,$NRInfo,"$out/$prefix.NR.txt","$out/$prefix.NR.stat.xls");

sub LoadTaxInfo{
	# 读取TaxInfo.txt.gz
	my %TaxName=();
	my %Tax=();
	open I,"gunzip -c $TaxInfo|";
	while(<I>){
		chomp;
		my @a=split/\t/;
		$TaxName{$a[1]}=$a[0]; # 物种名称到TaxID的映射
		$TaxName{$a[0]}=$a[1]; # TaxID到物种名，不再另开哈希了，偷个懒
		$Tax{$a[0]}=$a[-1]; # TaxID对应的层级信息
	}
	close I;
	return(\%TaxName,\%Tax); #返回多个哈希或数组务必先引用再解引用
}

sub GetFullTaxID(){
	# 前面读取了物种详细信息和指定的TaxID，那么接下来，提取该物种名下所有的TaxID
	my $TaxID=shift; 
	my $Tax=shift;
	my $TaxName=shift;
	my %Tax=%{$Tax}; #哈希解引用
	my %TaxName=%{$TaxName};
	my %SelectedTax=();
	$SelectedTax{$TaxID}="$TaxName{$TaxID}\t$Tax{$TaxID}"; #指定的Taxonomy当然要选
	my $KeyInfo=(split/\-/,$Tax{$TaxID})[0]; #获取TaxID的层级信息
	$KeyInfo="\-$TaxID\:$KeyInfo"; #横线-一定要添加，否则，若TaxID较短，则较长的TaxID只要有相同结尾，就会被选中
	foreach (keys %Tax){
		if($Tax{$_}=~m/$KeyInfo/){ 
			$SelectedTax{$_}="$TaxName{$_}\t$Tax{$_}"; #补一下物种名称
		} #这里就体现出了层级表的作用了。指定高级的分类单位，选择下属的全部低级的分类
	}
	return %SelectedTax; 
}


sub DB2Tax{
	# UniProt或NR库序列到Taxonomy
	my $SelectedTax=shift;
	my $DB=shift;
	my $out=shift;
	my $outstat=shift;
	my %Stat=(); #蛋白数目统计结果
	my %SelectedTax=%{$SelectedTax}; #解引用
	
	open I,"gunzip -c $DB|";
	open O,">$out";
	while(<I>){
		chomp;
		my @a=split/\t/;
		if($SelectedTax{$a[1]}){
			print O "$_\n";
			$Stat{$a[1]}++;
		}
	}
	close O;
	close I;
	
	open O,"|sort -k 2nr >$outstat"; #调用管道命令，将物种按照序列数目降序排列
	foreach(sort keys %Stat){
		print O "$_\t$Stat{$_}\t$SelectedTax{$_}\n";
	}
	close O;
}
